import sys
import os
import markdown
from dotenv import load_dotenv
from openai import OpenAI
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QTextEdit, QPushButton
from PyQt5.QtCore import Qt, QThread, pyqtSignal

# 加载环境变量
load_dotenv()

# AIWorker类
class AIWorker(QThread):
    update_signal = pyqtSignal(str)

    def __init__(self, user_input):
        super().__init__()
        self.user_input = user_input
        
        # 初始化 OpenAI 客户端
        self.client = OpenAI(
            base_url=os.getenv('OPENAI_API_BASE'),
            api_key=os.getenv('OPENAI_API_KEY')
        )

    def run(self):
        try:
            # 使用 OpenAI 兼容接口进行对话
            stream = self.client.chat.completions.create(
                model=os.getenv('OPENAI_API_MODEL', 'deepseek-chat'),
                messages=[{'role': 'user', 'content': self.user_input}],
                stream=True
            )

            # 处理流式响应
            for chunk in stream:
                if chunk.choices[0].delta.content is not None:
                    content = chunk.choices[0].delta.content
                    self.update_signal.emit(content)
        except Exception as e:
            self.update_signal.emit(f"\n[Error] {str(e)}")


class ChatWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.init_ui()
        self.ai_thread = None
        # 修正 Markdown 转换器初始化
        self.md = markdown.Markdown(
            extensions=['fenced_code', 'tables', 'nl2br', 'codehilite']
        )
        self.current_response = ""  # 用于累积流式响应
        self.last_cursor_position = 0  # 添加光标位置追踪

    def init_ui(self):
        self.setWindowTitle("InChat")
        self.setGeometry(100, 100, 600, 500)  # 设置窗口初始大小与位置

        layout = QVBoxLayout()  # 创建垂直布局

        # 修改聊天显示框以支持HTML
        self.chat_display = QTextEdit()
        self.chat_display.setReadOnly(True)
        self.chat_display.setAcceptRichText(True)
        self.chat_display.setStyleSheet("""
            QTextEdit {
                background-color: #f0f0f0;
                font-size: 14px;
                padding: 10px;
            }
            QTextEdit code {
                background-color: #e0e0e0;
                padding: 2px 4px;
                border-radius: 3px;
                font-family: monospace;
            }
        """)

        # 创建输入框，用户输入消息
        self.input_area = QTextEdit()
        self.input_area.setMaximumHeight(100)
        self.input_area.setPlaceholderText("输入您的问题...\n(Ctrl+Enter发送)")
        self.input_area.installEventFilter(self)  # 安装事件过滤器
        self.input_area.setStyleSheet("""
            background-color: white;
            font-size: 14px;
            padding: 10px;
        """)

        # 创建发送按钮，点击时发送消息
        self.send_btn = QPushButton("提问")
        self.send_btn.setStyleSheet("""
            QPushButton {
                background-color: #4CAF50;
                color: white;
                border: none;
                padding: 10px;
                font-size: 14px;
            }
            QPushButton:hover { background-color: #45a049; }
        """)
        self.send_btn.clicked.connect(self.send_message)  # 点击按钮时调用send_message

        # 将组件添加到布局中
        layout.addWidget(self.chat_display)
        layout.addWidget(self.input_area)
        layout.addWidget(self.send_btn)
        self.setLayout(layout)

    def eventFilter(self, obj, event):
        if obj == self.input_area and event.type() == event.KeyPress:
            if event.key() == Qt.Key_Return and event.modifiers() == Qt.ControlModifier:
                self.send_message()
                return True
        return super().eventFilter(obj, event)

    def send_message(self):
        # 获取用户输入的文本并去掉空白
        user_input = self.input_area.toPlainText().strip()
        if not user_input:
            return  # 如果输入为空，返回

        # 在聊天窗口中显示用户输入的消息
        self._append_message("You", user_input)
        self.input_area.clear()  # 清空输入框

        # 在聊天窗口中显示“思考中...”提示
        self._append_message("Bot", "思考中...", is_streaming=True)

        # 启动后台线程，处理与AI的对话
        self.ai_thread = AIWorker(user_input)
        # 连接信号和槽函数，更新聊天内容
        self.ai_thread.update_signal.connect(self.update_bot_response)
        # 连接线程完成后的槽函数
        self.ai_thread.finished.connect(self.finalize_response)
        # 启动线程
        self.ai_thread.start()

    def _append_message(self, sender, message, is_streaming=False):
        cursor = self.chat_display.textCursor()
        cursor.movePosition(cursor.End)

        if sender == "You":
            prefix = "\nYou: "
            message_html = f'<span style="color:#333333">{message}</span>'
        else:
            prefix = "\nBot: "
            if is_streaming:
                message_html = f'<span style="color:#888888">{message}</span>'
            else:
                # 将Markdown转换为HTML
                try:
                    # 重置转换器状态
                    self.md.reset()
                    message_html = self.md.convert(message)
                    message_html = f'<div style="color:#2196F3">{message_html}</div>'
                except:
                    message_html = f'<span style="color:#2196F3">{message}</span>'

        cursor.insertText(prefix)
        cursor.insertHtml(message_html)
        self.chat_display.ensureCursorVisible()

    def update_bot_response(self, content):
        self.current_response += content
        current_text = self.chat_display.toPlainText()
        cursor = self.chat_display.textCursor()
        
        if current_text.endswith("思考中..."):
            # 第一次响应，替换"思考中..."
            cursor.movePosition(cursor.End)
            for _ in range(7):  # "思考中..." 的长度
                cursor.deletePreviousChar()
            # 转换并显示新的响应
            try:
                self.md.reset()
                response_html = self.md.convert(self.current_response)
                cursor.insertHtml(f'<div style="color:#2196F3">{response_html}</div>')
                self.last_cursor_position = cursor.position()
            except:
                cursor.insertHtml(f'<span style="color:#2196F3">{self.current_response}</span>')
                self.last_cursor_position = cursor.position()
        else:
            # 获取新内容的HTML
            try:
                self.md.reset()
                # 只处理新增的内容
                response_html = self.md.convert(content)
                # 将光标移动到上次的位置
                cursor.setPosition(self.last_cursor_position)
                cursor.insertHtml(f'<span style="color:#2196F3">{response_html}</span>')
                self.last_cursor_position = cursor.position()
            except:
                cursor.setPosition(self.last_cursor_position)
                cursor.insertText(content)
                self.last_cursor_position = cursor.position()
                
        self.chat_display.setTextCursor(cursor)
        self.chat_display.ensureCursorVisible()

    def finalize_response(self):
        self.current_response = ""  # 清空累积的响应
        current = self.chat_display.toPlainText()
        if not current.endswith("\n"):
            self.chat_display.append("")

# 主程序入口
if __name__ == "__main__":
    app = QApplication(sys.argv)  # 创建应用程序
    window = ChatWindow()  # 创建聊天窗口实例
    window.show()  # 显示窗口
    sys.exit(app.exec_())  # 进入应用的主循环